=begin

  This file is part of MAMBO, a low-overhead dynamic binary modification tool:
      https://github.com/beehive-lab/mambo

  Copyright 2013-2016 Cosmin Gorgovan <cosmin at linux-geek dot org>

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.

=end

# Quick and dirty generator for emit-style function encoding wrappers for plugins
# Takes as argument the path to the C file of an instruction encoder generated by PIE

def get_filecode()
  "__EMIT_#{ARGV[0].gsub(/[^\w]/, "_").upcase}__"
end

def generate_header()
  puts "#ifdef PLUGINS_NEW"
  if (@header_only)
    puts "#ifndef #{get_filecode()}"
    puts "#define #{get_filecode()}"
  end
  puts "#include \"../dbm.h\"\n"
  puts "#include \"../#{ARGV[0].gsub(".c", ".h")}\""
  puts "#include \"plugin_support.h\"\n"
end

def generate_footer()
  if (@header_only)
    puts "#endif"
  end
  puts "#endif"
end

def rewrite_name(name)
  name.gsub('void ', 'void emit_')
end

def generate_body(name, fields, size)
  puts ")\n{"
  print "\t#{name}((uint#{@min_size * 8}_t **)(&ctx->code.write_p)"

  fields.each do |field|
    print ", #{field}"
  end

  puts ");"
  puts "\tctx->code.write_p += #{size};\n}"
end

def process_file(filename)
  fields = []
  name = ""
  size = 0
  File.readlines(filename).each do |line|
    if (line.match(/^void/))
      puts rewrite_name(line)
      line = line.split(' ')
      name = line[1]
      fields = []
    elsif (line.match(/address,?$/))
      print "\tmambo_context *ctx"
      #get the base instruction size, should be either uint16_t or uint32_t
      size = line.match(/[0-9]+/)[0].to_i / 8
      @min_size = size if (size < @min_size)
    elsif (line.match(/^\tunsigned int/))
      # function arguments, stripped of commas and excluding the address
      line = line.match(/[\w\s]+/)[0]
      print ",\n#{line}" 
      line = line.split(' ')
      fields.push(line[line.size-1])
    elsif (line.include?("**address = "))
      puts if (fields.size == 0)
      if (@header_only)
        puts ");"
      else
        size *= 2 if (line.include?(">>"))
        generate_body(name, fields, size)
      end
    end
  end
end

abort "Syntax: generate_emit_wrapper.rb <PIE_ENCODER.c> [header]" unless (ARGV[0])
@header_only = true if (ARGV[1] and ARGV[1] == "header")
@min_size = 4;

generate_header()
process_file(ARGV[0])
generate_footer()

